#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
修复训练数据提示词文件
1. 重命名提示词文件以匹配图片文件名
2. 优化提示词内容，以yuxin为核心关键词
"""

import os
import shutil
from pathlib import Path

def get_optimized_prompt(original_prompt, category):
    """根据原始提示词和类别生成优化的提示词"""
    
    # 基础关键词
    base_keywords = "yuxin, woman, female"
    
    # 根据类别添加特定关键词
    if "head" in category:
        return f"{base_keywords}, head shot, portrait, face, hair, updo, bun, side view"
    elif "hand" in category:
        return f"{base_keywords}, hand, fingers, palm, close up, side angle, pale skin, relaxed pose"
    elif "body" in category:
        return f"{base_keywords}, full body, dance pose, white sports bra, leggings, sneakers, dynamic pose, athletic wear"
    else:
        return f"{base_keywords}, female figure"

def fix_training_data():
    """修复训练数据"""
    
    base_path = Path("training_data_all/images")
    
    # 遍历所有子目录
    for subdir in base_path.iterdir():
        if not subdir.is_dir():
            continue
            
        print(f"处理目录: {subdir.name}")
        
        # 获取所有图片文件
        png_files = list(subdir.glob("*.png"))
        txt_files = list(subdir.glob("*.txt"))
        
        print(f"  找到 {len(png_files)} 个图片文件")
        print(f"  找到 {len(txt_files)} 个提示词文件")
        
        # 处理图片文件名，删除末尾的下划线
        for png_file in png_files:
            if png_file.stem.endswith('_'):
                new_name = png_file.stem[:-1] + png_file.suffix
                new_path = subdir / new_name
                png_file.rename(new_path)
                print(f"  重命名图片: {png_file.name} → {new_name}")
        
        # 重新获取处理后的图片文件列表
        png_files = list(subdir.glob("*.png"))
        
        # 创建新的提示词文件
        for i, png_file in enumerate(png_files):
            # 生成对应的提示词文件名
            txt_filename = png_file.stem + ".txt"
            txt_path = subdir / txt_filename
            
            # 如果已经有对应的提示词文件，读取内容
            if txt_path.exists():
                with open(txt_path, 'r', encoding='utf-8') as f:
                    original_prompt = f.read().strip()
                    # 删除包含3d的内容
                    if "3d" in original_prompt.lower():
                        original_prompt = "The image presents a female figure."
            else:
                # 使用默认提示词
                original_prompt = "The image presents a female figure."
            
            # 生成优化的提示词
            optimized_prompt = get_optimized_prompt(original_prompt, subdir.name)
            
            # 写入新的提示词文件
            with open(txt_path, 'w', encoding='utf-8') as f:
                f.write(optimized_prompt)
            
            print(f"  创建/更新: {txt_filename}")
        
        # 删除旧的提示词文件（如果存在）
        for txt_file in txt_files:
            if not any(txt_file.stem == png_file.stem for png_file in png_files):
                txt_file.unlink()
                print(f"  删除旧文件: {txt_file.name}")
        
        # 删除所有npz文件
        npz_files = list(subdir.glob("*.npz"))
        for npz_file in npz_files:
            npz_file.unlink()
            print(f"  删除npz文件: {npz_file.name}")

if __name__ == "__main__":
    print("开始修复训练数据...")
    fix_training_data()
    print("修复完成！") 